// Copyright (c) [Year] [name of copyright holder]
// [Software Name] is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2. 
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2 
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.  
// See the Mulan PSL v2 for more details.

/*!
 * \autor captainpeng
 * \date 2019-5-15
 * \version 1.0
 */

#ifndef MY_AVLTREE_HPP
#define MY_AVLTREE_HPP

#include"mystd.hpp"
#include"function.hpp"

namespace my{

    template<typename T>
    class avlNode{
	public:
		using elementType = T;
		using heightType = short;

		elementType mValue;
        avlNode * mLeft;
        avlNode * mRight;
        heightType mHeight;

		avlNode():mValue(), mLeft(nullptr), mRight(nullptr), mHeight(0){}

        void swapNodeInfo(avlNode & rhs){
            // 交换高度
            heightType htmp = mHeight;
            mHeight = rhs.mHeight;
            rhs.mHeight = htmp;

            // 交换左右子树
            avlNode * ttmp = mLeft;
            mLeft = rhs.mLeft;
            rhs.mLeft = ttmp;

            ttmp = mRight;
            mRight = rhs.mRight;
            rhs.mRight = ttmp;
        }
	};
    
    template<typename T>
    class avl{
    public:
		using nodeType = avlNode<info_type_t<T>>;
        using elementType = typename nodeType::elementType;
        using heightType = typename nodeType::heightType;

		using treeType = nodeType *;
        using memType = my::memory;

        static const int AVL_MAX_HEIGHT = 1024;

	private:
        treeType mHead;
        size_t mSize;

        static heightType Height(treeType node){
            if(node == nullptr) return -1;
            return node->mHeight;
        }
        
        static void HeightCal(treeType node){
            if(node == nullptr) return ;
            node->mHeight = MY_MAX(Height(node->mLeft), Height(node->mRight)) + 1;
        }
        
        static treeType SingleRotateWithLeft(treeType k2){
            treeType k1;

            k1 = k2->mLeft;
            k2->mLeft = k1->mRight;
            k1->mRight = k2;

            HeightCal(k1);
            HeightCal(k2);

            return k1;
        }

        static treeType SingleRotateWithRight(treeType k2){
            treeType k1;

            k1 = k2->mRight;
            k2->mRight = k1->mLeft;
            k1->mLeft = k2;

            HeightCal(k1);
            HeightCal(k2);

            return k1;
        }
        
        static treeType DoubleRotateWithLeft(treeType k3){
            k3->mLeft = SingleRotateWithRight(k3->mLeft);

            return SingleRotateWithLeft(k3);
        }

        static treeType DoubleRotateWithRight(treeType k3){
            k3->mRight = SingleRotateWithLeft(k3->mRight);

            return SingleRotateWithRight(k3);
        }

        static void * AllocDebug(size_t size){
            return memFunc<nodeType>::Alloc(size);
        }

        static void FreeDebug(void * ptr){
            return memFunc<nodeType>::Free(ptr);
        }

        inline static treeType Alloc(size_t size){
#ifdef MY_ENABLE_DEBUG
            treeType cur = memType::calloc<nodeType>(1, AllocDebug, FreeDebug);
#else
            treeType cur = memType::calloc<nodeType>(1);
#endif
            return cur;
        }

        template<typename COMP, typename ... ARGS>
        static treeType Find(treeType cur, treeType * S, char * V, int & now, const COMP & comp, ARGS && ... args){
            while(cur != nullptr){
				if(++now >= AVL_MAX_HEIGHT)
                    return nullptr;

				S[now] = cur;
                int result = comp(cur->mValue, std::forward<ARGS>(args)...);

                if(result > 0){
                    V[now] = 'l';
                    cur = cur->mLeft;
                }else if(result < 0){
                    V[now] = 'r';
                    cur = cur->mRight;
                }else break;
            }

            return cur;
        }

        static treeType BalanceTree(treeType cur, treeType * S, char * V, int & now){
            for(;now >= 0; --now){
                if(V[now] == 'l'){
                    S[now]->mLeft = cur;
                }else if(V[now] == 'r'){
                    S[now]->mRight = cur;
                }

				cur = S[now];
				if(Height(cur->mLeft) - Height(cur->mRight) == 2)
					if(V[now+1] == 'l')
						cur = SingleRotateWithLeft(cur);
					else
						cur = DoubleRotateWithLeft(cur);
				else if(Height(cur->mRight) - Height(cur->mLeft) == 2)
					if(V[now+1] == 'r')
						cur = SingleRotateWithRight(cur);
					else
						cur = DoubleRotateWithRight(cur);

                HeightCal(cur);
            }

            return cur;
        }

        template<typename COMP, typename ... ARGS>
        static treeType EraseFind(treeType cur, treeType & del, treeType * S, char * V, int & now, const COMP & comp, ARGS && ... args){
            treeType temp = nullptr;
            // 1. 先查找，栈记录着查找的路径，找到合适位置后再进行旋转平衡操作
            while(cur != nullptr){
                
                int result = comp(cur->mValue, std::forward<ARGS>(args)...);

				if(result > 0){
                    // result > 0 -> cur->mValue > el
                    S[++now] = cur;
                    V[now] = 'l';
                    cur = cur->mLeft;
                }else if(result < 0){
                    // result < 0 -> cur->mValue < el
                    S[++now] = cur;
                    V[now] = 'r';
                    cur = cur->mRight;
                }else if(cur->mLeft != nullptr && cur->mRight != nullptr){
                    S[++now] = cur;
                    // 拥有两个子节点
                    V[now] = 'r';
                    int curnow = now;

                    // temp = EraseFindMin(cur->mRight, S, V, now);
                    // cur->mValue = std::move((*temp).mValue);
                    temp = cur->mRight;
                    while(temp->mLeft != nullptr){
                        S[++now] = temp;
                        V[now]  = 'l';
                        temp = temp->mLeft;
                    }

                    // 交换节点信息，不交换值的信息
                    // 交换之后
                    // cur->节点的信息 == 原 temp 的，但 cur->值的信息 == 原 cur 的
                    // temp->节点的信息 == 原 cur 的，但 temp->值的信息 == 原 temp 的
                    // 那么 cur 将变成要被删掉的，而 temp 变为要保存到 curnow 的

                    cur->swapNodeInfo(*temp);
                    S[curnow] = temp;
                    temp = cur;

                    // 不管 temp->mRight 有没有子树，cur 都要保存下来
                    cur = temp->mRight;
                    temp->mRight = nullptr;

                    break;
                }else{
                    // 拥有一个或者零个子节点
                    temp = cur;
                    if(cur->mLeft == nullptr)
                        cur = cur->mRight;
                    else
                        cur = cur->mLeft;

                    break;
                }
            }

            del = temp;
            return cur;
        }

        static treeType EraseBalanceTree(treeType cur, treeType * S, char * V, int & now){
            // 注意，这里与 BalanceTree 函数不一样，因为该平衡树的判断规则和函数不一样
            for(;now >= 0; --now){
                if(V[now] == 'l'){
                    S[now]->mLeft = cur;
                }else if(V[now] == 'r'){
                    S[now]->mRight = cur;
                }

                cur = S[now];

                if(Height(cur->mLeft) - Height(cur->mRight) == 2){
					if(cur->mLeft->mLeft != nullptr)
						cur = SingleRotateWithLeft(cur);
					else
						cur = DoubleRotateWithLeft(cur);
                }else if(Height(cur->mRight) - Height(cur->mLeft) == 2)
					if(cur->mRight->mRight != nullptr)
						cur = SingleRotateWithRight(cur);
					else
						cur = DoubleRotateWithRight(cur);

                HeightCal(cur);
            }

            return cur;
        }

        static void PrintTreeNode(std::ostream & os, treeType cur, char s, int count){
            for(int i = 0; i < count; ++i)
                os<<"  ";

			if(cur == nullptr){
				os<<"NULL"<<std::endl;
				return;
			}else
				os<<"["<<s<<", "<<cur->mHeight<<": "<<cur->mValue<<"]"<<std::endl;
			
			PrintTreeNode(os, cur->mLeft, 'l', count+1);
			PrintTreeNode(os, cur->mRight, 'r', count+1);
        }

    public:
        avl():mHead(nullptr){}

        avl(avl && rhs){
            mHead = rhs.mHead;
            rhs.mHead = nullptr;
        }

        avl & operator=(avl && rhs){
            free();
            mHead = rhs.mHead;
            rhs.mHead = nullptr;            
        }

        ~avl(){
            free();
        }

        size_t size() const {
            return mSize;
        }

        elementType * insert(elementType && el){
            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;

			// treeType cur = mHead;
            // while(true){
            //     if(cur == nullptr){
            //         cur = Alloc(1);
			// 		// std::cout<<"cur == "<<(void *)cur<<std::endl;
            //         cur->mValue = std::move(el);
            //         break;
            //     }

			// 	if(++now >= AVL_MAX_HEIGHT){
            //         std::cerr<<"my::avl::insert 已超出树最大的高度，不能插入"<<std::endl;
            //         return nullptr;
            //     }

			// 	S[now] = cur;
            //     int result = func<elementType>::comp(cur->mValue, el);

            //     if(result > 0){
            //         // cur->mValue > el
            //         V[now] = 'l';
			// 		// std::cout<<"["<<now<<"] left"<<std::endl;
            //         cur = cur->mLeft;
            //     }else if(result < 0){
            //         // cur->mValue < el
            //         V[now] = 'r';
			// 		// std::cout<<"["<<now<<"] right"<<std::endl;
            //         cur = cur->mRight;
            //     }else{
            //         cur->mValue = std::move(el);
			// 		return &cur->mValue;
            //     }
            // }
            // mHead = BalanceTree(cur, S, V, now);
            // return &cur->mValue;

            treeType cur = Find(mHead, S, V, now, func<elementType>::comp, el);
            if(cur == nullptr){
                if(now >= AVL_MAX_HEIGHT){
                    std::cerr<<"my::avl::insert, 已超出树最大的高度，不能插入"<<std::endl;
                    return nullptr;
                }

                cur = Alloc(1);
                ++mSize;

                mHead = BalanceTree(cur, S, V, now);
            }

            cur->mValue = std::move(el);
            return &cur->mValue;
        }
        
        elementType * insert(const elementType & el){
            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;

			// treeType cur = mHead;
            // while(true){
            //     if(cur == nullptr){
            //         cur = Alloc(1);
			// 		// std::cout<<"cur == "<<(void *)cur<<std::endl;
            //         cur->mValue = el;
            //         break;
            //     }
				
			// 	if(++now >= AVL_MAX_HEIGHT){
            //         std::cerr<<"my::avl::insert 已超出树最大的高度，不能插入"<<std::endl;
            //         return nullptr;
            //     }

			// 	S[now] = cur;
            //     int result = func<elementType>::comp(cur->mValue, el);

            //     if(result > 0){
            //         // cur->mValue > el
            //         V[now] = 'l';
			// 		// std::cout<<"["<<now<<"] left"<<std::endl;
            //         cur = cur->mLeft;
            //     }else if(result < 0){
            //         // cur->mValue < el
            //         V[now] = 'r';
			// 		// std::cout<<"["<<now<<"] right"<<std::endl;
            //         cur = cur->mRight;
            //     }else{
            //         cur->mValue = el;
			// 		return &cur->mValue;
            //     }
            // }
            // mHead = BalanceTree(cur, S, V, now);
            // return &cur->mValue;

            treeType cur = Find(mHead, S, V, now, func<elementType>::comp, el);
            if(cur == nullptr){
                if(now >= AVL_MAX_HEIGHT){
                    std::cerr<<"my::avl::insert, 已超出树最大的高度，不能插入"<<std::endl;
                    return nullptr;
                }

                cur = Alloc(1);
                ++mSize;

                mHead = BalanceTree(cur, S, V, now);
            }

            cur->mValue = el;

            return &cur->mValue;
        }

        // 初始化，条件插入
        template<typename INIT, typename COMP, typename ... ARGS>
        elementType * insertComp(const INIT & init, const COMP & comp, ARGS && ... args){
            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;

			treeType cur = Find(mHead, S, V, now, comp, std::forward<ARGS>(args)...);

            // while(cur != nullptr){
			// 	if(++now >= AVL_MAX_HEIGHT){
            //         std::cerr<<"my::avl::insertComp 已超出树最大的高度，不能插入"<<std::endl;
            //         return nullptr;
            //     }

			// 	S[now] = cur;

            //     // 注意， cond 为当前值和输入值进行比较的结果
            //     // cond > 0 , 当前值大于输入值
            //     // cond < 0 , 当前值小于当前值
            //     int cond = comp(cur->mValue, std::forward<ARGS>(args)...);
			// 	if(cond > 0){
            //         V[now] = 'l';
			// 		// std::cout<<"["<<now<<"] left"<<std::endl;
            //         cur = cur->mLeft;
            //     }else if(cond < 0){
            //         V[now] = 'r';
			// 		// std::cout<<"["<<now<<"] right"<<std::endl;
            //         cur = cur->mRight;
            //     }else break;
            // }

            if(cur == nullptr){
                cur = Alloc(1);
                ++mSize;
                mHead = BalanceTree(cur, S, V, now);
            }

            // std::cout<<"cur == "<<(void *)cur<<std::endl;
            init(cur->mValue, std::forward<ARGS>(args)...);

            return &cur->mValue;
        }

        template<typename COMP, typename ... ARGS>
        elementType * findComp(const COMP & comp, ARGS && ... args){
            treeType cur = mHead;
            while(true){
                if(cur == nullptr) return nullptr;
                // 注意， cond 为当前值和输入值进行比较的结果
                // cond > 0 , 当前值大于输入值
                // cond < 0 , 当前值小于当前值
                int cond = comp(cur->mValue, std::forward<ARGS>(args)...);
                if(cond > 0){
                    cur = cur->mLeft;
                }else if(cond < 0){
                    cur = cur->mRight;
                }else
                    return &cur->mValue;
            }
        }

        elementType * find(const elementType & el){
            treeType cur = mHead;
            while(true){
                if(cur == nullptr) return nullptr;

                if(el < cur->mValue)
                    cur = cur->mLeft;
                else if(el > cur->mValue)
                    cur = cur->mRight;
                else
                    return &cur->mValue;
            }
        }

        elementType * findMin(){
            treeType cur = mHead;
            while(cur->mLeft != nullptr){
                cur = cur->mLeft;
            }

            return &cur->mValue;
        }

        elementType * findMax(){
            treeType cur = mHead;
            while(cur->mRight != nullptr){
                cur = cur->mRight;
            }

            return &cur->mValue;
        }

		void free(){
            if(mHead == nullptr) return;

            treeType cur = mHead;
            treeType temp = nullptr;

            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = 0;

            while(cur != nullptr){
                if(cur->mLeft != nullptr && cur->mRight != nullptr){
                    // 拥有两个子节点
                    S[++now] = cur;
                    V[now] = 'r';
                    int curnow = now;

                    // temp = EraseFindMin(cur->mRight, S, V, now);
                    // cur->mValue = std::move((*temp).mValue);
                    temp = cur->mRight;
                    while(temp->mLeft != nullptr){
                        S[++now] = temp;
                        V[now]  = 'l';
                        temp = temp->mLeft;
                    }

                    // 交换节点信息，不交换值的信息
                    // 交换之后
                    // cur->节点的信息 == 原 temp 的，但 cur->值的信息 == 原 cur 的
                    // temp->节点的信息 == 原 cur 的，但 temp->值的信息 == 原 temp 的
                    // 那么 cur 将变成要被删掉的，而 temp 变为要保存到 curnow 的

                    cur->swapNodeInfo(*temp);
                    S[curnow] = temp;
                    temp = cur;

                    // 不管 temp->mRight 有没有子树，cur 都要保存下来
                    cur = temp->mRight;
                    temp->mRight = nullptr;

                    break;
                }else{
                    // 拥有一个或者零个子节点
                    temp = cur;
                    if(cur->mLeft == nullptr)
                        cur = cur->mRight;
                    else
                        cur = cur->mLeft;

                    break;
                }

                cur = EraseBalanceTree(cur, S, V, now);
                memType::free(temp);
            }

			mHead = nullptr;
            mSize = 0;
		}

        void erase(const elementType & el){
            if(mHead == nullptr) return ;

            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;


            treeType del = nullptr;
            treeType cur = EraseFind(mHead, del, S, V, now, func<elementType>::comp, el);

            if(del != nullptr){
                mHead = EraseBalanceTree(cur, S, V, now);
                memType::free(del);
                --mSize;
            }
        }

        template<typename COMP, typename ... ARGS>
        void eraseComp(const COMP & comp, ARGS && ... args){
            if(mHead == nullptr) return ;

            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;


            treeType del = nullptr;
            treeType cur = EraseFind(mHead, del, S, V, now, comp, std::forward<ARGS>(args) ... );

            if(del != nullptr){
                mHead = EraseBalanceTree(cur, S, V, now);
                memType::free(del);
                --mSize;
            }

        }

        template<typename CALL, typename ... ARGS>
        void each(const CALL & call, ARGS && ... args){
            treeType S[AVL_MAX_HEIGHT];
            char V[AVL_MAX_HEIGHT];
            int now = -1;

            treeType cur = mHead;
            bool back_flag = false;

            while(cur != nullptr){
                // std::cout<<cur<<std::endl;
                // ... 别问我为啥写出个这个
                // back_flag 记录着 cur 是否是回溯的
                // 如果 cur->mLeft 为 null，那 continue 就不会执行。
                // 如果 cur->mRight 为 null，那 continue 和 改变back_flag 都不会执行
                // 最后则是调用输入的 call ，那里就可以改变回溯的 back_flag
                if(!back_flag){
                    S[++now] = cur;
                    V[now] = 'l';
                    if(cur->mLeft != nullptr){
                        cur = cur->mLeft;
                        continue;
                    }
                }

                if(V[now] == 'l'){
                    V[now] = 'r';
                    if(cur->mRight != nullptr){
                        cur = cur->mRight;
                        back_flag = false;
                        continue;
                    }
                }

                if(V[now] == 'r'){
                    // 回调函数的返回一个 bool 值告诉是否还要继续遍历
                    if(!call(cur->mValue, std::forward<ARGS>(args)...))
                        return ;
                    cur = (now > 0)? S[--now]:nullptr;
                    back_flag = true;
                }
            }
        }

        void print(std::ostream & os){
            PrintTreeNode(os, mHead, 'h', 0);
        }
    };
}

#endif
